import threading
import datetime
from typing import List, Dict, Any
from database.milvus_client import MilvusClient
from pymilvus import (
    FieldSchema,
    DataType
)

import logging
from models.chat_record import ChatRecord, ContentType, Role
from server.gme_inference import GmeQwen2VL
import numpy as np

logger = logging.getLogger(__name__)


class ImageTextVacServer(MilvusClient):
    _instance = None
    _lock = threading.Lock()
    # 类变量定义集合字段（避免重复初始化）
    collection_name = "image_text_retrieval_s"
    collection_fields = [
        FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
        FieldSchema(name="user_id", dtype=DataType.VARCHAR, max_length=36),
        FieldSchema(name="saas_id", dtype=DataType.VARCHAR, max_length=36),
        FieldSchema(name="image_path", dtype=DataType.VARCHAR, max_length=512),
        FieldSchema(name="text_desc", dtype=DataType.VARCHAR, max_length=1024),
        FieldSchema(name="image_vec", dtype=DataType.FLOAT_VECTOR, dim=1536),
        FieldSchema(name="text_vec", dtype=DataType.FLOAT_VECTOR, dim=1536)
    ]

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            with cls._lock:
                if not cls._instance:
                    cls._instance = super().__new__(cls)
                    cls._instance._initialized = False  # 初始化标志位
        return cls._instance

    def __init__(self, host="localhost", port="19530", user="root", pwd="szsk8888"):
        if self._initialized:  # 已初始化则直接返回
            return

        # 模型初始化（仅执行一次）
        print(f"加载多模态模型: gme-Qwen2-VL-2B-Instruct")
        # self.gme = GmeQwen2VL("D:/model/qwenvl2b/iic/gme-Qwen2-VL-2B-Instruct")
        self.gme =None
        print(f"成功加载多模态模型")

        # 父类初始化
        super().__init__(
            host=host,
            port=port,
            collection_name=self.collection_name,
            fields=self.collection_fields,
            user=user,
            pwd=pwd
        )

        self._initialized = True  # 标记已初始化

    def create_index(self):
        """创建向量索引"""
        # 创建高效索引
        index_params = {
            "index_type": "IVF_PQ",  # 平衡精度和速度
            "metric_type": "IP",  # CLIP使用内积相似度
            "params": {
                "nlist": 256,
                "m": 32,  # 压缩维度
                "nbits": 8
            }
        }
        self.collection.create_index("image_vec", index_params)
        self.collection.create_index("text_vec", index_params)

    def encode_image(self,image_path):
        """用CLIP模型生成图像特征向量"""

        # input_img = load_image(image_path)  # 支持皮卡丘示例图片路径/本地图片 返回PIL.Image
        # # 支持一张图片(PIL.Image)或多张图片(List[PIL.Image])输入，输出归一化特征向量
        # return pipeline.forward({'img': input_img})['img_embedding'].numpy().astype(np.float64).flatten()  # 2D Tensor, [图片数, 特征维度]
        imgs=[image_path]

        return self.gme.get_image_embeddings(images=imgs).numpy().astype(np.float64).flatten()
        #return "111"
    def encode_text(self,text):
        """用文本模型生成特征向量"""
        # 支持一条文本(str)或多条文本(List[str])输入，输出归一化特征向量
        # return pipeline.forward({'text': text})['text_embedding'].numpy().astype(np.float64).flatten()  # 2D Tensor, [文本数, 特征维度]
        texts=[text]
        return self.gme.get_text_embeddings(texts=texts).numpy().astype(np.float64).flatten()
        #return "22"

    def insert_record(self,  image_path, description,user_id,saas_id) -> str:
        """插入新记录"""
        try:
            # 生成双向量
            img_vec = self.encode_image(image_path)
            print(type(img_vec))
            text_vec = self.encode_text(description)
            # 插入数据
            data = [
                [user_id],
                [saas_id],
                [image_path],  # image_url
                [description],  # description
                [img_vec],  # image_embedding
                [text_vec]  # text_embedding
            ]
            print("开始："+datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            mr = self.collection.insert(data)
            # self.collection.flush()
            print("结束：" + datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            return "1111111"
        except Exception as e:
            logger.error(f"插入记录失败: {str(e)}")
            raise RuntimeError(f"插入记录失败: {str(e)}") from e

    def search_records(
            self,
            user_id: str,
            saas_id: str,
            query: str,
            alpha=0.8,
            top_k: int = 1
    ) -> List[Dict[str, Any]]:
        """
        混合检索（支持文本/图像输入）
        :param alpha: 图像特征权重（0-1），text权重=1-alpha
        """
        # 判断输入类型
        if isinstance(query, str):  # 文本查询
            query_text_vec = self.encode_text(query)
            query_img_vec = self.encode_text(query)  # 用文本描述模拟图像查询
        else:  # 图像路径/PIL Image
            query_img_vec = self.encode_image(query)
            query_text_vec = self.encode_image(query)  # 用图像特征模拟文本查询

        # 双向量搜索
        search_params = {"metric_type": "IP", "params": {"nprobe": 10}}
        # 构建查询表达式
        expr_parts = f'user_id == "{user_id}"'
        expr_parts.append(f'saas_id == "{saas_id}"')

        # 并行搜索
        img_results = self.collection.search(
            [query_img_vec], "image_vec", search_params, limit=top_k * 2, expr = expr_parts,
            output_fields=["image_path", "text_desc"]
        )

        text_results = self.collection.search(
            [query_text_vec], "text_vec", search_params, limit=top_k * 2, expr = expr_parts,
            output_fields=["image_path", "text_desc"]
        )

        # 结果融合
        fused_results = []
        for hits in img_results:
            for hit in hits:
                print(f"img_vec_score；{hit.score}")
                fused_results.append({
                    "id": hit.id,
                    "score": hit.score * alpha,  # 图像分数加权
                    "type": "image",
                    **{k: hit.entity.get(k) for k in ["image_path", "text_desc"]}
                })

        for hits in text_results:
            for hit in hits:
                exists = next((r for r in fused_results if r["id"] == hit.id), None)
                if exists:
                    exists["score"] += hit.score * (1 - alpha)  # 累加文本分数
                else:
                    fused_results.append({
                        "id": hit.id,
                        "score": hit.score * (1 - alpha),
                        "type": "text",
                        **{k: hit.entity.get(k) for k in ["image_path", "text_desc"]}
                    })

        # 按总分排序
        return sorted(fused_results, key=lambda x: -x["score"])[:top_k]